from launch import LaunchDescription
from launch_ros.actions import Node
from launch.actions import DeclareLaunchArgument
from launch.substitutions import LaunchConfiguration

def generate_launch_description():
    # 声明启动参数
    model_type_arg = DeclareLaunchArgument(
        'model_type',
        default_value='bert',
        description='模型类型: bert 或 vit'
    )
    
    device_arg = DeclareLaunchArgument(
        'device',
        default_value='cuda:0',
        description='运行设备: cuda:0, cpu 等'
    )
    
    partitions_arg = DeclareLaunchArgument(
        'partitions',
        default_value='2',
        description='分区数量'
    )
    
    # 获取参数值
    model_type = LaunchConfiguration('model_type')
    device = LaunchConfiguration('device')
    partitions = LaunchConfiguration('partitions')
    
    # 定义节点
    master_node = Node(
        package='voltage_ros2',
        executable='master_node',
        name='voltage_master',
        parameters=[
            {'model_type': model_type},
            {'device': device},
            {'partitions': partitions}
        ],
        output='screen'
    )
    
    worker_node = Node(
        package='voltage_ros2',
        executable='worker_node',
        name='voltage_worker',
        parameters=[
            {'worker_id': 1},
            {'device': device},
            {'partitions': partitions}
        ],
        output='screen'
    )
    
    inference_service = Node(
        package='voltage_ros2',
        executable='inference_service',
        name='voltage_inference',
        parameters=[
            {'model_type': model_type},
            {'device': device}
        ],
        output='screen'
    )
    
    # 创建启动描述
    return LaunchDescription([
        model_type_arg,
        device_arg,
        partitions_arg,
        master_node,
        worker_node,
        inference_service
    ])